package com.tc.sharding;

import com.google.common.collect.Range;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;
import org.springframework.util.CollectionUtils;

import java.util.*;

public class DataTableShardingAlgorithm implements ComplexKeysShardingAlgorithm {
    @Override
    public Collection<String> doSharding(Collection collection, ComplexKeysShardingValue complexKeysShardingValue) {
        List<String> ret = new ArrayList<>();

        Map columnNameAndShardingValuesMap = complexKeysShardingValue.getColumnNameAndShardingValuesMap();
        Map columnNameAndRangeValuesMap = complexKeysShardingValue.getColumnNameAndRangeValuesMap();

        List<Long> timeList = (List<Long>) columnNameAndShardingValuesMap.get("time_stamp");
        if (!CollectionUtils.isEmpty(timeList)){
            ret.add("batch_data_" + getYear(timeList.get(0)) + "_" + getMonth(timeList.get(0)));
            return ret;
        }


        Range<Long> range = (Range<Long>) columnNameAndRangeValuesMap.get("time_stamp");
        Long upperRange = range.upperEndpoint();
        Long lowerRange = range.lowerEndpoint();

        for (int i = getYear(lowerRange); i <= getYear(upperRange); i++){
            int k = 0;
            if (i < getYear(upperRange)){
                k = 12;
            }else{
                k = getMonth(upperRange);
            }
            for (int j = getMonth(lowerRange); j <= k; j++){
                ret.add("batch_data_" + i + "_" + j);
            }
        }
//        List<String> month = LocalDateUtils.listDateStrsNew("2024-07", "2024-08", "month");
        return ret;
    }

    private Integer getYear(Long timestamp){
        Calendar calendar = Calendar.getInstance();
        calendar.setTimeInMillis(timestamp);

        int year = calendar.get(Calendar.YEAR);
        return year;
    }

    private Integer getMonth(Long timestamp){
        Calendar calendar = Calendar.getInstance();
        calendar.setTimeInMillis(timestamp);

        int month = calendar.get(Calendar.MONTH) + 1;
        return month;
    }

    private String getYearStr(Long timestamp){
        Calendar calendar = Calendar.getInstance();
        calendar.setTimeInMillis(timestamp);

        int year = calendar.get(Calendar.YEAR);
        return year + "";
    }

    private String getMonthStr(Long timestamp){
        Calendar calendar = Calendar.getInstance();
        calendar.setTimeInMillis(timestamp);

        int month = calendar.get(Calendar.MONTH) + 1;
        if (month < 10){
            return "0" + month;
        }
        return month + "";
    }
}
